#!/usr/bin/python

# Discrete Event Simulator
# Richard Skowyra - rskowyra@bu.edu
# William Blair - wdblair@bu.edu

RANDOM_SEED = 2233 #The random seed for numpy
SIMTIME = 100000
DATA_PATH = '../node-data/'
NODE_POPULATION = 100
TOTAL_ITEMS = 1000
UNIFORM_QUERY_RATE = 10 #Inter-arrival times
UNIFORM_NOISE_P = 0.01
NUM_BASIS_VECTORS = 2
MEAN_VERSIONS = 3
MEAN_VERSION_TTL = 100
DATA_COLLECTION_INTERVAL = 1000
MODE = 1 #0=uniform query function, 1=bps, 2=Exp3

from heapq import *
import numpy
import random
import pickle

#Seed the random number generator.
numpy.random.seed(RANDOM_SEED)

#Make a uniform sample of size k * n on average where k*n <= n if fixed = False, otherwise, it is fixed.
def sample(list,k,fixed=False):
    list_size = len(list)
    sample = []
    if not fixed:
        k = numpy.random.binomial(list_size,k,1)[0]
    
    for i in xrange(k):
        sample.append(list[i])

    for j in xrange(k,list_size):
        r = numpy.random.random_integers(0,j,1)[0]
        if r < k:
            sample[r] = list[j]
    return sample

class Event:
    def __init__(self,iat):
        self.iat = iat #Inter-arrival time
    
    def next_arrival(self):
        return numpy.random.exponential(self.iat,1)[0]
        
    def __le__(self,other):  #Define the <= operator for events (used by the heap implementation)
        return self.timestamp <= other.timestamp
    
    def run(self,controller):
        pass
    
    def schedule(self,timestamp):
        self.timestamp = timestamp

class Query(Event):
    def __init__(self, iat, node):
        Event.__init__(self,iat)
        self.node = node
    
    def run(self,controller):
        self.node.query(controller)
        next_query = Query(self.iat,self.node)
        controller.schedule_event(next_query)

class TTLExpire(Event):
    def __init__(self,iat,item,version):
        Event.__init__(self,iat)
        self.item = item
        self.version = version

    def run(self,controller):
        for k,v in controller.BV.items():
            controller.BV[k][self.item][self.version] = numpy.random.random_integers(0,1,1)[0]
            
        for node in controller.nodes:
            if self.item in node.versions:
                node.expire_version(controller,self.item,self.version)
        
        next_expiration = TTLExpire(self.iat,self.item,self.version)
        controller.schedule_event(next_expiration)

class Snapshot(Event):
    def __init__(self,iat):
        Event.__init__(self,iat)
        
    def run(self,controller):
        for node in controller.nodes:
            with open(node.path, 'a') as f:
                pickle.dump(node.costs, f)
            node.costs = []
        
        controller.schedule_event( Snapshot(self.iat) )

class Node:
    def __init__(self,node_id, bv_id,contacts, noise):
        self.node_id = node_id
        self.noise = noise
        self.bv_id = bv_id
        self.contacts = contacts
        self.versions = {} #Dictionary of the form [item][ver][pref]
        self.costs = []
        self.path = DATA_PATH + str(self.node_id) + '.dat'

    def expire_version(self, controller, item, version):
        if version in self.versions[item]: del self.versions[item][version]
    
    def add_version(self, controller, item,version):
        lp = (random.random < self.noise) ^ controller.BV[self.bv_id][item][version]
        try:
            self.versions[item][version] = lp
        except KeyError:
            self.versions[item] = {}
            self.versions[item][version] = lp                    
    
    def recommend(self, item, version_set):
        if item in self.versions:
            recs = [x for x in self.versions[item].keys() if (x in version_set and self.versions[item][x] == 1)]
            if recs: rec = random.choice(recs)
            else: return False
        else:
            return False

    ###Uniform Selection with No Recommendations (maximum cost)
    def query_uniform(self, controller, item, version_set):
        cost = 0
        random.shuffle(version_set)
        for version in version_set:
            lp = (random.random < self.noise) ^ controller.BV[self.bv_id][item][version]
            try:
                self.versions[item][version] = lp
            except KeyError:
                self.versions[item] = {}
                self.versions[item][version] = lp
            if lp: break                    
            cost+=1
        return cost
    
    ###Selection Function used by Boaz Patt-Shamir
    def query_bps(self, controller, item, version_set):
        cost = -1
        noise_factor = random.random < self.noise
        lp = False
        while((not lp) and version_set):
            cost += 1
            #Do we ask for a recommendation from someone?
            if random.random() < 0.5: choose_rec = 0
            else: choose_rec = 1
            if choose_rec == 1: #Ask for a recommendation
                contact = random.choice(self.contacts)
                version = contact.recommend(item, version_set)
                if not version: 
                    cost -= 1 #No cost assigned for asking someone who does not know
                    continue

            else: #Do not ask for a recommendation
                version = random.choice(version_set)
                
            lp = noise_factor ^ controller.BV[self.bv_id][item][version]
            try:
                self.versions[item][version] = lp
            except KeyError:
                self.versions[item] = {}
                self.versions[item][version] = lp
            version_set.remove(version)
            
        return cost
                
    def query(self,controller):
        mode = MODE
        #Choose an with no known good version
        item = random.choice(controller.versions.keys())
        while item in self.versions: #Already have this item
            if 1 in self.versions[item].values(): #Have a good version, re-choose
                item = random.choice(controller.versions.keys())
            else: #No good versions, item choice is fine
                break
        #Restrict selection to unknown versions
        version_set = controller.BV[self.bv_id][item].keys()
        if item in self.versions:
            version_set = [x for x in version_set if x not in self.versions[item]]
        
        #Query the item
        if mode == 0:
            self.costs.append(self.query_uniform(controller, item, version_set))
        elif mode==1:
            self.costs.append(self.query_bps(controller, item, version_set))
        else:
            raise ValueError('Invalid mode')
    
class Controller:
    def __init__(self):
        self.time = 0.0
        self.schedule = []
        self.BV = {}
        self.nodes = []
        self.items = []
        self.versions = {}
        self.basis_cnt = 0
    
    def create_items(self,count):
        self.items = [x for x in xrange(count)]
        
    def create_nodes(self,query_rates, noise_dist):
        #Create basis vector assignments
        #Random Assignment
        #bvs = numpy.random.random_integers(0,self.basis_cnt-1,len(query_rates))
        #Ordered assignment
        bvs = numpy.array([x % self.basis_cnt for x in xrange(len(query_rates))])
        node_ids = range(len(query_rates))
        for idx, q in enumerate(query_rates):
            node_id = node_ids[idx]
            bv_id = bvs[idx]
            noise = noise_dist[idx]
            new_node = Node(node_id,bv_id,[], noise)
            self.nodes.append(new_node)
            first_query = Query(q,new_node)
            self.schedule_event(first_query)
        for node in self.nodes:   #By default, make all nodes know of all other nodes.
            node.contacts = self.nodes

    #Make some random versions for each item.
    def create_versions(self, versions_lambda, ttl_lambda):
        #For each item, create versions
        for i in xrange(len(self.items)):
            item = self.items[i]
            nversions = versions_lambda #Number of versions. Constant number right now.
            self.versions[item] = range(nversions) #Monotonically increasing version ids
            for j in xrange(nversions): #For each version
                ttl = numpy.random.exponential(ttl_lambda)
                expire = TTLExpire(ttl, item, j)
                self.schedule_event(expire) #Schedule its first expiration.

                
    #Initialize the basis vectors
    def setup_bv(self):
        sample = numpy.random.random_integers(0,1,self.basis_cnt*len(self.items)) #Create a sample
        for i in xrange(self.basis_cnt):
	    self.BV[i] = {}
            for idx,j in enumerate(self.items):
		self.BV[i][j] = {}
                sample = numpy.random.random_integers(0,1,len(self.versions[j]))
                for l,k in enumerate(self.versions[j]):
                    self.BV[i][j][k] = sample[l]
  
    def setup_node_versions(self,nodes_with_versions,items_per_node,versions_per_item): #Populate some initial item/version pairs to nodes.
        nodes_to_samples = sample(self.nodes,nodes_with_versions)
        for node in self.nodes:
            items_to_add = sample(self.items,items_per_node)
            for item in items_to_add:
                versions_to_add = sample(self.versions[item],versions_per_item)
                for version in versions_to_add:
                    node.add_version(self, item,version)
                    
    def schedule_event(self,event):
        event.schedule( self.time + event.next_arrival() )
        heappush(self.schedule,event)
        
    def run(self):
        while (self.time < SIMTIME):
            print self.time
            try:
                event = heappop(self.schedule)
                self.time = event.timestamp
                event.run(self)
            except IndexError: #We've run out of events
                break
            
    def setup_data_collection(self,iat):
        #Create and/or truncate the output files as needed.
        for node in self.nodes:
            with open(node.path,'w') as f:
                pass
        self.schedule_event(Snapshot(iat))

if __name__ == "__main__":
    print 'Initializing random state'
    #Create node query rate list
    #Uniform
    query_rates = [UNIFORM_QUERY_RATE]*NODE_POPULATION
    #Create node noise list
    noise_dist = [UNIFORM_NOISE_P]*NODE_POPULATION
    #Create node and item state
    controller = Controller()
    controller.basis_cnt = NUM_BASIS_VECTORS
    print 'Creating nodes'
    controller.create_nodes(query_rates, noise_dist)  #10000 Nodes with 5 to 25 queries per unit time. Replace with any other list, where len=#nodes and items are query rates
    print 'Creating items'
    controller.create_items(TOTAL_ITEMS)
    print 'Creating versions'
    controller.create_versions(MEAN_VERSIONS, MEAN_VERSION_TTL) #Create, on average, 25 versions per item with an average of .25 TTL expirations per unit time. #Change to E[TTL]
    #Setup the simulator's state
    print 'Initializing simulation state'
    controller.setup_bv()
    controller.setup_node_versions(0.9,0.001,0.20) #On average, 90% of nodes have versions for .1% of the global items, and 20% of the versions for each item
    #Setup the data collection event
    print 'Setting Up Data Collection'
    controller.setup_data_collection(DATA_COLLECTION_INTERVAL)
    print 'Running simulation'
    controller.run()
    print 'Simulation complete.'
